from implement.layers.basic.model import Model
from implement.models.decoder import Decoder
from implement.models.encoder import Encoder
from utils.functions_collect import binary_cross_entropy, flatten, exp, sum

'''
Variational Autoencoder（VAE）是一种生成模型，用于学习数据的潜在表示。它属于自编码器（Autoencoder）的一种扩展，具有潜在变量（latent variables）。VAE的主要目标是学习一个能够生成输入数据的潜在表示，同时保持对数据的合理解释。
具体来说，VAE包括两个主要部分：
编码器（Encoder）： 将输入数据映射到潜在空间中的潜在表示。编码器输出潜在变量的均值和方差。
解码器（Decoder）： 将潜在表示映射回数据空间，生成与输入数据尽可能接近的重建。
在训练过程中，VAE通过最大化似然下界（ELBO）来学习模型参数。ELBO包括两部分：重建损失（reconstruction loss）和正则化项，通常是KL散度（Kullback-Leibler divergence），用于确保学到的潜在表示与标准正态分布接近。
VAE的主要应用领域包括生成新样本、数据压缩、以及学习数据的连续潜在表示，例如在图像生成、语音处理和自然语言处理等任务中取得了良好的效果。
'''
class VAE(Model):
    def __init__(self, latent_size):
        """Variational Autoencoder (VAE) 构造函数。

        Args:
            latent_size (int): 潜在空间的维度。
        """
        super().__init__()
        self.encoder = Encoder(latent_size)
        self.decoder = Decoder()

    def forward(self, x, C=1.0, k=1):
        """计算 VAE 的损失函数，即 Evidence Lower Bound (ELBO) 乘以 -1。

        Args:
            x (Variable or ndarray): 输入变量。
            C (int): 通常为 1.0。可以更改以控制 ELBO 边界的第二项，该项充当正则化。
            k (int): 在编码向量中使用的蒙特卡罗采样数量。
        """
        z_mean, z_log_var = self.encoder(x)

        rec_loss = 0
        for l in range(k):
            z = self.encoder.sampling(z_mean, z_log_var)
            y = self.decoder(z)
            rec_loss += binary_cross_entropy(flatten(y), flatten(x)) / k

        kl_loss = C * (z_mean ** 2 + exp(z_log_var) - z_log_var - 1) * 0.5
        kl_loss = sum(kl_loss) / len(x)
        return rec_loss + kl_loss

